package com.hps.lock;

import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.locks.InterProcessMutex;
import org.apache.curator.retry.ExponentialBackoffRetry;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

/**
 * @author heps
 * @date 2019/1/21 15:42
 */
public class ZkUtil {

    private InterProcessMutex mutex;
    private CuratorFramework client;

    public boolean lock() {
        RetryPolicy retryPolicy = new ExponentialBackoffRetry(200, 3);
        client = CuratorFrameworkFactory.newClient("127.0.0.1:2181", retryPolicy);
        client.start();
        mutex = new InterProcessMutex(client, "/curator/lock");
        try {
            boolean acquire = mutex.acquire(1000, TimeUnit.MILLISECONDS);
            if (acquire) {
                return true;
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return false;
    }

    public void release() {
        try {
            mutex.release();
            client.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public void close() {
        client.close();
    }

    public static void main(String[] args) {
        ExecutorService executorService = Executors.newFixedThreadPool(30);
        List<Integer> list = new ArrayList<>();
        list.add(0);
        int taskNum = 1000;
        CountDownLatch latch = new CountDownLatch(taskNum);
        long start = System.currentTimeMillis();
        for (int i = 0; i < taskNum; i++) {
            executorService.submit(() -> {
                ZkUtil zkUtil = new ZkUtil();
                boolean result = zkUtil.lock();
                if (result) {
                    Integer num = list.get(list.size() - 1);
                    list.add(++num);
                    zkUtil.release();
                }
                zkUtil.close();
                latch.countDown();
            });
        }
        try {
            latch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        System.out.println(System.currentTimeMillis() - start);
        executorService.shutdown();
        System.out.println(list.size());
        System.out.println(list);
    }
}
